package sshx

import (
	"bytes"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/zeromicro/go-zero/core/logx"
	"golang.org/x/crypto/ssh"
)

// SSHClient encapsulates SSH connection and operations
type SSHClient struct {
	client *ssh.Client
	config *ssh.ClientConfig
	server SSHConfig
}

// NewSSHClient creates a new SSH connection
func NewSSHClient(config SSHConfig) (*SSHClient, error) {
	// Validate parameters
	if err := ValidateSSHConfig(&config); err != nil {
		return nil, err
	}

	// Create SSH connection configuration
	var authMethods []ssh.AuthMethod

	// Use key or password for authentication
	if config.KeyPath != "" {
		key, err := os.ReadFile(config.KeyPath)
		if err != nil {
			return nil, fmt.Errorf("failed to read SSH key file: %v", err)
		}

		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("failed to parse SSH private key: %v", err)
		}
		authMethods = append(authMethods, ssh.PublicKeys(signer))
	} else if config.Password != "" {
		authMethods = append(authMethods, ssh.Password(config.Password))
	} else {
		return nil, fmt.Errorf("must provide password or key file")
	}

	// Set up SSH client configuration
	sshConfig := &ssh.ClientConfig{
		User:            config.User,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // In production environment, use ssh.FixedHostKey()
		Timeout:         30 * time.Second,
	}

	// Connect to SSH server
	addr := fmt.Sprintf("%s:%d", config.Host, config.Port)
	client, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		return nil, fmt.Errorf("SSH connection failed: %v", err)
	}

	return &SSHClient{
		client: client,
		config: sshConfig,
		server: config,
	}, nil
}

// Close closes the SSH connection
func (s *SSHClient) Close() error {
	if s.client != nil {
		return s.client.Close()
	}
	return nil
}

// ExecuteCommand executes a command on the SSH server
func (s *SSHClient) ExecuteCommand(command string) (string, error) {
	// Create a new session
	session, err := s.client.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// Execute command and get output
	output, err := session.CombinedOutput(command)
	if err != nil {
		return string(output), fmt.Errorf("command execution failed: %v, output: %s", err, output)
	}

	return string(output), nil
}

// ExecuteCommandWithOutput executes a command on the SSH server and returns stdout and stderr separately
func (s *SSHClient) ExecuteCommandWithOutput(command string) (stdout, stderr string, err error) {
	// Create a new session
	session, err := s.client.NewSession()
	if err != nil {
		return "", "", fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// Set up stdout and stderr buffers
	var stdoutBuf, stderrBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	session.Stderr = &stderrBuf

	// Execute command
	err = session.Run(command)
	return stdoutBuf.String(), stderrBuf.String(), err
}

// ExecuteRemoteCommand connects to another server through the current server and executes a command
func (s *SSHClient) ExecuteRemoteCommand(targetConfig SSHConfig, command string) (string, error) {
	// Validate parameters
	if err := ValidateSSHConfig(&targetConfig); err != nil {
		return "", err
	}

	// Build command based on authentication method
	var remoteCmd string
	if targetConfig.Password != "" {
		// Handle special characters in password
		escapedPassword := strings.Replace(targetConfig.Password, "'", "'\\''", -1)
		remoteCmd = fmt.Sprintf("sshpass -p '%s' ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 -p %d %s@%s '%s'",
			escapedPassword, targetConfig.Port, targetConfig.User, targetConfig.Host, command)
	} else if targetConfig.KeyPath != "" {
		// Use key authentication
		remoteCmd = fmt.Sprintf("ssh -i %s -o StrictHostKeyChecking=no -o ConnectTimeout=5 -p %d %s@%s '%s'",
			targetConfig.KeyPath, targetConfig.Port, targetConfig.User, targetConfig.Host, command)
	} else {
		return "", fmt.Errorf("no password or key provided, cannot connect to target server")
	}

	// Execute remote command
	return s.ExecuteCommand(remoteCmd)
}

// CheckSshpassInstalled checks if sshpass is installed
func (s *SSHClient) CheckSshpassInstalled() (bool, error) {
	output, err := s.ExecuteCommand("which sshpass")
	return err == nil && len(output) > 0, nil
}

// InstallSshpass installs sshpass
func (s *SSHClient) InstallSshpass() (bool, error) {
	// Detect Linux distribution and install sshpass
	installCmd := `
if command -v apt-get >/dev/null 2>&1; then
    sudo apt-get update -qq && sudo apt-get install -y sshpass
elif command -v yum >/dev/null 2>&1; then
    sudo yum install -y sshpass
elif command -v dnf >/dev/null 2>&1; then
    sudo dnf install -y sshpass
elif command -v zypper >/dev/null 2>&1; then
    sudo zypper install -y sshpass
elif command -v apk >/dev/null 2>&1; then
    sudo apk add sshpass
else
    echo "Unrecognized package manager, cannot automatically install sshpass"
    exit 1
fi
`
	output, err := s.ExecuteCommand(installCmd)
	if err != nil {
		logx.Errorf("sshpass installation failed: %v, output: %s", err, output)
		return false, fmt.Errorf("sshpass installation failed: %v", err)
	}

	// Check again if sshpass is installed
	return s.CheckSshpassInstalled()
}

// EnsureRemoteDirectory ensures the directory exists on the remote server
func (s *SSHClient) EnsureRemoteDirectory(targetConfig SSHConfig, remoteDir string) error {
	// Check if the directory exists
	exists, err := s.CheckFileExists(remoteDir)
	if err != nil {
		return fmt.Errorf("failed to check if directory exists: %v", err)
	}
	if exists {
		return nil
	}

	// Build mkdir command
	mkdirCmd := fmt.Sprintf("mkdir -p \"%s\"", remoteDir)

	// Execute mkdir command on the target server
	output, err := s.ExecuteRemoteCommand(targetConfig, mkdirCmd)
	if err != nil {
		return fmt.Errorf("failed to create target directory: %v, output: %s", err, output)
	}

	logx.Infof("Target directory created successfully: %s", remoteDir)
	return nil
}

// CheckFileExists checks if a file exists and is readable
func (s *SSHClient) CheckFileExists(filePath string) (bool, error) {
	// Build check command
	checkCmd := fmt.Sprintf("[ -f \"%s\" ] && [ -r \"%s\" ] && echo OK || echo NOT_FOUND",
		filePath, filePath)

	// Execute check command
	output, err := s.ExecuteCommand(checkCmd)
	if err != nil {
		return false, err
	}

	return strings.Contains(output, "OK"), nil
}

// GetFileInfo gets file size and MD5
func (s *SSHClient) GetFileInfo(filePath string) (size string, md5 string, err error) {
	// Get file size and MD5
	fileInfoCmd := fmt.Sprintf("stat -c %%s \"%s\" && md5sum \"%s\" | awk '{print $1}'",
		filePath, filePath)

	output, err := s.ExecuteCommand(fileInfoCmd)
	if err != nil {
		return "", "", fmt.Errorf("failed to get file info: %v, output: %s", err, output)
	}

	// Parse output
	lines := strings.Split(strings.TrimSpace(output), "\n")
	if len(lines) < 2 {
		return "", "", fmt.Errorf("failed to parse file info: %s", output)
	}

	size = strings.TrimSpace(lines[0])
	md5 = strings.TrimSpace(lines[1])

	return size, md5, nil
}

// CheckRemoteFile checks if a file exists on the remote server and gets its information
func (s *SSHClient) CheckRemoteFile(targetConfig SSHConfig, filePath string) (exists bool, size string, md5 string, err error) {
	// Build check command
	checkCmd := fmt.Sprintf("[ -f \"%s\" ] && echo \"EXISTS\" && stat -c %%s \"%s\" && md5sum \"%s\" | awk \"{print \\$1}\" || echo \"NOT_EXISTS\"",
		filePath, filePath, filePath)

	// Execute remote command
	output, err := s.ExecuteRemoteCommand(targetConfig, checkCmd)
	if err != nil {
		return false, "", "", err
	}

	// Parse output
	lines := strings.Split(strings.TrimSpace(output), "\n")
	if len(lines) < 1 || lines[0] != "EXISTS" {
		return false, "", "", nil
	}

	if len(lines) < 3 {
		return true, "", "", fmt.Errorf("failed to get complete file information")
	}

	return true, strings.TrimSpace(lines[1]), strings.TrimSpace(lines[2]), nil
}

// ScpFileTo uses SCP to transfer a file to a remote server
func (s *SSHClient) ScpFileTo(targetConfig SSHConfig, sourceFilePath, targetDir string) (stdout, stderr string, err error) {
	// Build target address
	targetAddr := fmt.Sprintf("%s@%s:%s",
		targetConfig.User, targetConfig.Host, targetDir)

	// Build SCP command
	scpCmd := fmt.Sprintf("scp -o StrictHostKeyChecking=no -o ConnectTimeout=10 -P %d \"%s\" %s",
		targetConfig.Port, sourceFilePath, targetAddr)

	// If the target server requires password authentication, use sshpass
	if targetConfig.Password != "" {
		// Handle special characters in password
		escapedPassword := strings.Replace(targetConfig.Password, "'", "'\\''", -1)
		scpCmd = fmt.Sprintf("sshpass -p '%s' %s", escapedPassword, scpCmd)
	}

	// Execute SCP command
	return s.ExecuteCommandWithOutput(scpCmd)
}
